#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#define UP_DIV(x, y) (((x) + (y) - (1)) / (y))

extern void Conv8x8(float *input, float *weight, float *output, const float *bias, size_t step, size_t offset);

typedef struct {
    int kernel_h_;
    int kernel_w_;
    int input_channel_;
    int output_h_;
    int output_w_;
    int output_channel_;
} ConvParameter;

void test(float *input_data, float *weight_data, float *output_data, const float *bias_data,
          ConvParameter *conv_param)
{
    int kernel_h = conv_param->kernel_h_;
    int kernel_w = conv_param->kernel_w_;
    int in_channel = conv_param->input_channel_;
    int out_h = conv_param->output_h_;
    int out_w = conv_param->output_w_;
    int out_channel = conv_param->output_channel_;

    float *gemm_input = input_data;
    float *gemm_weight = weight_data;
    float *gemm_output = output_data;

    int channel_block = UP_DIV(in_channel, 4);
    int kernel_plane = kernel_h * kernel_w;
    int unit_size = kernel_plane * channel_block * 4;
    // we take 8 conv blocks, to compute 8 channels of 8 outputs per loop
    int oc_batches = UP_DIV(out_channel, 8);
    int batches = UP_DIV(out_h * out_w, 8);
    int total_loops = oc_batches * batches;
    // we accumulate 4 channels per time for input blocks
    int ic_batches = UP_DIV(in_channel, 4);
    int conv_depth = kernel_h * kernel_w * ic_batches;
    // bytes from one output's i-th channel to the next output's i-th channel
    int output_offset = out_channel * sizeof(float);

    int i = 0;
    while (i++ < total_loops) {
        Conv8x8(gemm_input, gemm_weight, gemm_output, bias_data, conv_depth, output_offset);
        // please add offset instead of multiply by i
        gemm_input += 8 * unit_size;
        if (i % oc_batches != 0) {
            gemm_weight += unit_size;
            gemm_output += 8;
        } else {
            gemm_weight = weight_data;
            gemm_output += 7 * out_channel + 8;
        }
    }
}

int main()
{
    int kernel_h = 2;
    int kernel_w = 2;
    int in_channel = 4;
    int out_h = 8;
    int out_w = 1;
    int out_channel = 8;
    int tile_num = 8;
    int inchannel_block = 4;
    ConvParameter conv_param = {
        kernel_h,
        kernel_w,
        in_channel,
        out_h,
        out_w,
        out_channel
    };

    int out_plane = out_h * out_w;
    int channel_block = UP_DIV(in_channel, inchannel_block);
    int kernel_plane = kernel_h * kernel_w;
    int unit_size = kernel_plane * channel_block * inchannel_block;
    int packed_input_size = out_plane * unit_size;
    float *input_data = (float *) malloc(packed_input_size * sizeof(float));
    for (int i = 0; i < packed_input_size; i++) {
        input_data[i] = (float) (i / 32 + 1);
    }

    int kernel_block = UP_DIV(out_channel, tile_num);
    int kernel_size = kernel_plane * in_channel * out_channel;
    int pack_weight_size = kernel_block * channel_block * tile_num * inchannel_block * kernel_plane;
    float *weight_data = (float *) malloc(pack_weight_size * sizeof(float));
    for (int i = 0; i < pack_weight_size; i++) {
        weight_data[i] = (float) (i % 4 + 1);
    }

    float *bias_data = (float *) malloc(out_channel * sizeof(float));
    for (int i = 0; i < out_channel; i++) {
        bias_data[i] = (float) (i % out_channel);
    }

    float *output_data = (float *) malloc(out_plane * out_channel * sizeof(float));
    memset(output_data, 0, out_plane * out_channel * sizeof(float));

    test(input_data, weight_data, output_data, bias_data, &conv_param);

    float res[] = {40, 81, 122, 163, 44, 85, 126, 167};
    for (int i = 0; i < out_plane * out_channel; i++) {
        if (output_data[i] != res[i % 8]) {
            printf("index %d wrong\n", i);
        }
    }

    return 0;
}